import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader, random_split
import random

# 读取变换矩阵数据
def read_transformation_matrices(file_path):
    """
    从CSV文件中读取刚体变换矩阵
    
    参数:
        file_path (str): CSV文件路径
        
    返回:
        list: 包含每个变换矩阵的字典列表，每个字典包含'rotation'和'translation'键
    """
    df = pd.read_csv(file_path)
    
    matrices = []
    for index, row in df.iterrows():
        # 提取旋转矩阵部分
        rotation = np.array([
            [row['Vx0'], row['Vy0'], row['Vz0']],
            [row['Vx1'], row['Vy1'], row['Vz1']],
            [row['Vx2'], row['Vy2'], row['Vz2']]
        ])
        
        # 提取平移向量部分
        translation = np.array([row['P0'], row['P1'], row['P2']])
        
        matrices.append({
            'rotation': rotation,
            'translation': translation
        })
    
    return matrices

# 自定义数据集类
class TransformationDataset(Dataset):
    def __init__(self, A, B):
        self.A = A
        self.B = B
    
    def __len__(self):
        return len(self.A)
    
    def __getitem__(self, idx):
        # 将旋转矩阵和平移向量展平为一维数组
        a = np.concatenate([self.A[idx]['rotation'].flatten(), self.A[idx]['translation']])
        b = np.concatenate([self.B[idx]['rotation'].flatten(), self.B[idx]['translation']])
        
        return torch.FloatTensor(a), torch.FloatTensor(b)

# 定义神经网络模型
class TransformationModel(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(TransformationModel, self).__init__()
        self.fc1 = nn.Linear(input_dim, 128)
        self.fc2 = nn.Linear(128, 256)
        self.fc3 = nn.Linear(256, 128)
        self.fc4 = nn.Linear(128, output_dim)
        self.relu = nn.ReLU()
    
    def forward(self, x):
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        x = self.relu(self.fc3(x))
        x = self.fc4(x)
        return x

# 训练模型
def train_model(model, train_dataloader, criterion, optimizer, num_epochs=100):
    model.train()
    for epoch in range(num_epochs):
        running_loss = 0.0
        for inputs, targets in train_dataloader:
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            loss.backward()
            optimizer.step()
            running_loss += loss.item() * inputs.size(0)
        epoch_loss = running_loss / len(train_dataloader.dataset)
        print(f'Epoch {epoch+1}/{num_epochs}, Loss: {epoch_loss:.8f}')

# 测试模型
def test_model(model, test_dataloader, criterion):
    model.eval()
    test_loss = 0.0
    with torch.no_grad():
        for inputs, targets in test_dataloader:
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            test_loss += loss.item() * inputs.size(0)
    test_loss /= len(test_dataloader.dataset)
    print(f'Test Loss: {test_loss:.8f}')
    return test_loss

# 计算全局误差
def calculate_global_error(model, A, B):
    model.eval()
    total_error = 0.0
    with torch.no_grad():
        for a, b in zip(A, B):
            # 准备输入数据
            input_data = np.concatenate([a['rotation'].flatten(), a['translation']])
            input_tensor = torch.FloatTensor(input_data).unsqueeze(0)
            
            # 预测
            output = model(input_tensor)
            
            # 计算误差
            pred_rotation = output[0, :9].reshape(3, 3).numpy()
            pred_translation = output[0, 9:].numpy()
            
            rotation_diff = np.linalg.norm(pred_rotation - b['rotation'], 'fro')
            translation_diff = np.linalg.norm(pred_translation - b['translation'])
            
            total_error += rotation_diff + translation_diff
    
    return total_error

# 保存模型权重
def save_model_weights(model, file_path):
    torch.save(model.state_dict(), file_path)
    print(f"Model weights saved to {file_path}")

# 加载模型权重
def load_model_weights(model, file_path):
    model.load_state_dict(torch.load(file_path))
    print(f"Model weights loaded from {file_path}")
    return model

# 使用模型进行预测并输出结果
def predict_and_display(model, A, B, num_samples=5):
    model.eval()
    with torch.no_grad():
        for i in range(10):

            j = random.randint(0, len(A)-1)

            a = A[j]
            b = B[j]
            
            # 准备输入数据
            input_data = np.concatenate([a['rotation'].flatten(), a['translation']])
            input_tensor = torch.FloatTensor(input_data).unsqueeze(0)
            
            # 预测
            output = model(input_tensor)
            
            # 提取预测结果
            pred_rotation = output[0, :9].reshape(3, 3).numpy()
            pred_translation = output[0, 9:].numpy()
            
            # 打印结果
            print(f"Sample {j+1}:")
            print("Original A Matrix:")
            print("Rotation:")
            print(a['rotation'])
            print("Translation:")
            print(a['translation'])
            print("\nTransformed A Matrix:")
            print("Rotation:")
            print(pred_rotation)
            print("Translation:")
            print(pred_translation)
            print("\nCorresponding B Matrix:")
            print("Rotation:")
            print(b['rotation'])
            print("Translation:")
            print(b['translation'])
            print("\n" + "="*50 + "\n")

# 主函数
def main():
    # 替换为你的CSV文件路径
    file_path_A = 'foundation_pose_matrices1.csv'
    file_path_B = 'matrices.csv'
    
    # 读取数据
    A = read_transformation_matrices(file_path_A)
    B = read_transformation_matrices(file_path_B)
    
    # 创建数据集
    dataset = TransformationDataset(A, B)
    
    # 随机划分训练集和测试集
    train_size = 300
    test_size = len(dataset) - train_size
    train_dataset, test_dataset = random_split(dataset, [train_size, test_size])
    
    # 创建数据加载器
    train_dataloader = DataLoader(train_dataset, batch_size=32, shuffle=True)
    test_dataloader = DataLoader(test_dataset, batch_size=32, shuffle=False)
    
    # 创建模型
    input_dim = 3 * 3 + 3  # 旋转矩阵3x3 + 平移向量3
    output_dim = 3 * 3 + 3  # 旋转矩阵3x3 + 平移向量3
    model = TransformationModel(input_dim, output_dim)
    
    # 定义损失函数和优化器
    criterion = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=0.0001)
    
    # 训练模型
    train_model(model, train_dataloader, criterion, optimizer, num_epochs=1000)
    
    # 保存模型权重
    save_model_weights(model, 'transformation_model_weights.pth')
    
    # 测试模型
    test_loss = test_model(model, test_dataloader, criterion)
    
    # 计算全局误差
    global_error = calculate_global_error(model, A, B)
    print(f'Global Error: {global_error:.8f}')
    
    # 输出具体转换结果
    predict_and_display(model, A, B, num_samples=5)

if __name__ == "__main__":
    main()